from langchain_core.prompts import FewShotPromptTemplate, PromptTemplate
from src.module.Model import Model

examples = [
    {
        'question': '你的版本是多少',
        'answer': '版本为 7b llm'
    },
    {
        'question': '你的名字是什么',
        'answer': '我的名字叫小新'
    },
    {
        'question': '小皮球是什么颜色的',
        'answer': '小皮球是蓝色的'
    },
    {
        'question': '今天天气如何',
        'answer': '晴天'
    }
]

class FewShort(Model):
    def start(self):

        ## 短的格式化输出内容
        # 将问题与答案进行格式化输出
        prompt = FewShotPromptTemplate(
            examples=examples, ## 示例，对于一段列表示例
            example_prompt=PromptTemplate.from_template('问题: {question} \n {answer}'), ## 将示例内容转换为格式化内容
            suffix="问题：{input}", # 将转换的格式化内容，与问题进行结合转换为 问题加示例的格式化内容
        ).invoke({ "input": '小皮球是什么颜色的? 今天天气如何？' }).to_string()

        print(self.model.invoke(prompt))


